package org.wenshu.ai;

import dev.langchain4j.model.openai.OpenAiChatModel;
import dev.langchain4j.store.embedding.EmbeddingMatch;
import org.junit.jupiter.api.Test;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.function.Function;

import com.fasterxml.jackson.databind.ObjectMapper;
import com.fasterxml.jackson.core.JsonProcessingException;
import dev.langchain4j.data.embedding.Embedding;
import dev.langchain4j.model.embedding.EmbeddingModel;
import dev.langchain4j.model.jina.JinaEmbeddingModel;
import dev.langchain4j.store.embedding.inmemory.InMemoryEmbeddingStore;
import dev.langchain4j.data.segment.TextSegment;
import dev.langchain4j.store.embedding.EmbeddingStore;
import dev.langchain4j.data.message.UserMessage;
import dev.langchain4j.data.message.SystemMessage;
import dev.langchain4j.data.message.ChatMessage;

import java.util.Arrays;

import static java.time.Duration.ofSeconds;

/**
 * Java实现的Agent示例
 * @author AI Generated
 */
class JavaSimpleAgentIT {

    // 配置OpenAI客户端
    private final OpenAiChatModel model = OpenAiChatModel.builder()
            .baseUrl("https://api.deepseek.com")
            .apiKey("your-key")
            .modelName("deepseek-chat")
            .logRequests(true)
            .logResponses(true)
            .build();

    // 配置嵌入模型
    private final EmbeddingModel embeddingModel = JinaEmbeddingModel.builder()
            .apiKey("your-key")
            .modelName("jina-embeddings-v3")
            .timeout(ofSeconds(10))
            .logRequests(true)
            .logResponses(true)
            .build();

    // 内存存储
    private final List<Map.Entry<String, String>> shortMem = new ArrayList<>(); // 保存最近的对话
    private final EmbeddingStore<TextSegment> longMem = new InMemoryEmbeddingStore<>(); // 存储所有向量
    private final ObjectMapper objectMapper = new ObjectMapper();

    // 工具定义
    private final Map<String, Function<String, Object>> tools = new HashMap<>();

    public JavaSimpleAgentIT() {
        // 初始化工具
        tools.put("calculator", this::calculator);
        tools.put("search", this::search);
    }

    /**
     * 计算器工具
     */
    private Object calculator(String expression) {
        try {
            // 简单的表达式计算，实际应用中可能需要更安全的实现
            javax.script.ScriptEngineManager manager = new javax.script.ScriptEngineManager();
            javax.script.ScriptEngine engine = manager.getEngineByName("JavaScript");
            return engine.eval(expression);
        } catch (Exception e) {
            return "表达式计算错误: " + e.getMessage();
        }
    }

    /**
     * 搜索工具
     */
    private String search(String query) {
        // 这里应该实现实际的搜索功能，目前返回模拟数据
        return "搜索结果: " + query + "。这是一个模拟的搜索结果。";
    }

    /**
     * 构建提示信息
     */
    private String buildPrompt(String question) {
        StringBuilder prompt = new StringBuilder("最近的对话记录:\n");
        
        // 添加最近的对话历史
        int historyLimit = Math.min(shortMem.size(), 3);
        for (int i = shortMem.size() - historyLimit; i < shortMem.size(); i++) {
            Map.Entry<String, String> entry = shortMem.get(i);
            prompt.append("Q:").append(entry.getKey()).append("\n");
            prompt.append("A:").append(entry.getValue()).append("\n");
        }
        
        // 添加相关记忆（如果有的话）
        try {
            // 将问题转换为嵌入向量
            Embedding questionEmbedding = embeddingModel.embed(question).content();

            // 查找最相关的记忆
            List<EmbeddingMatch<TextSegment>> relevantMemories = longMem.findRelevant(questionEmbedding, 1);
            if (!relevantMemories.isEmpty()) {
                prompt.append("\n相关记忆:\n").append(relevantMemories.get(0).embedded().text()).append("\n");
            }
        } catch (Exception e) {
            // 如果嵌入过程出错，忽略相关记忆部分
            System.err.println("查找相关记忆时出错: " + e.getMessage());
        }

        // 添加工具调用指导
        prompt.append("""
        请按步骤进行推理：
        1. 思考阶段：仔细思考用户的问题
        2. 工具调用阶段：选择可以调用的工具，并输出相应工具所需的参数
        如果你认为需要使用工具（如'calculator'计算器或'search'搜索），请以JSON对象的形式回应：
        {
            "tool": "<工具名称>",
            "params": "<参数>"
        }
        <工具名称>只能是calculator或search。不要生成其他内容。直接返回json，而不是带```json```块。
        例如，如果你需要计算，回应类似这样：
        {"tool": "calculator", "params": "3 + 4"}
        """);

        // 添加新问题
        prompt.append("\n新问题: ").append(question);

        return prompt.toString();
    }

    /**
     * 获取API响应
     */
    private String getApiResponse(String prompt) {
        List<ChatMessage> messages = Arrays.asList(
            new SystemMessage("我是一个乐于助人的智能助手。"),
            new UserMessage(prompt)
        );

        return model.generate(messages).content().text();
    }

    /**
     * 更新记忆
     */
    private void updateMemories(String question, String response) {
        // 更新短期记忆
        shortMem.add(Map.entry(question, response));
        if (shortMem.size() > 3) {
            shortMem.remove(0);
        }
        
        // 更新长期记忆
        String memoryText = "Q:" + question + " A:" + response;
        TextSegment segment = TextSegment.from(memoryText);
        longMem.add(embeddingModel.embed(memoryText).content(), segment);
    }

    /**
     * 提取工具请求
     */
    private Map<String, String> extractToolRequest(String response) {
        try {
            // 尝试解析JSON响应
            return objectMapper.readValue(response, Map.class);
        } catch (JsonProcessingException e) {
            // 如果不是有效的JSON，则返回null
            return null;
        }
    }

    /**
     * 调用工具
     */
    private Object callTool(String toolName, String params) {
        if (tools.containsKey(toolName)) {
            return tools.get(toolName).apply(params);
        }
        return null;
    }

    /**
     * 主要的Agent函数
     */
    private String agent(String question) {
        // 构建提示
        String prompt = buildPrompt(question);
        
        // 获取LLM响应
        String response = getApiResponse(prompt);
        System.out.println("LLM response: " + response);
        
        // 提取工具请求
        Map<String, String> toolRequest = extractToolRequest(response);
        
        // 如果有工具请求，则调用工具
        if (toolRequest != null && toolRequest.containsKey("tool") && toolRequest.containsKey("params")) {
            String toolName = toolRequest.get("tool");
            String params = toolRequest.get("params");
            
            Object toolResult = callTool(toolName, params);
            
            // 使用工具结果生成最终响应
            response = getApiResponse("根据问题整理工具执行结果：" + toolResult + "。问题是：" + question + "。");
        }
        
        // 更新记忆
        updateMemories(question, response);
        
        return response;
    }

    @Test
    void testAgent() {
        // 测试Agent
        String question1 = "25乘以16等于多少？";
        System.out.println("问题: " + question1);
        String answer1 = agent(question1);
        System.out.println("Answer: " + answer1);
        
        String question2 = "现任美国总统是谁？";
        System.out.println("\n问题: " + question2);
        String answer2 = agent(question2);
        System.out.println("Answer: " + answer2);
    }

    /**
     * 交互式测试（可选）
     */
    public static void main(String[] args) {
        JavaSimpleAgentIT agentIT = new JavaSimpleAgentIT();
        java.util.Scanner scanner = new java.util.Scanner(System.in);
        
        while (true) {
            System.out.print("你: ");
            String question = scanner.nextLine();
            
            if ("quit".equalsIgnoreCase(question)) {
                break;
            }
            
            String answer = agentIT.agent(question);
            System.out.println("机器人: " + answer);
        }
        
        scanner.close();
    }
}